# chapter 4 加入校验的种子生成器自动生成代码

# chapter 3中，由于没有对代码的校验机制，导致有时候会执行失败。
# 本章节加入校验和修改机制，提升种子生成成功概率。


import ast
import os

# 基本环境

from openai import OpenAI
key = "" #换成你的key
url = "" # 换成你的url
model_name = "" # 换成你的模型名，建议使用qwen3-32b以上的模型
def print_llm(text):
    print(text, end="", flush=True)
def run_llm(prompt_str):
    result = ""
    client = OpenAI(api_key=key, base_url=url)
    response = client.chat.completions.create(
        model=model_name,
                    messages=[
                        {"role": "user", "content": prompt_str},
                    ],
                    stream=True,
                    # 把timeout值设置高一些
                    timeout=600
                )
    for chunk in response:
        content = chunk.choices[0].delta.content
        if content is not None:
            result += content
            print_llm(content)
    return result

wrapper_code = """


target_floder = "./j40/corpus"
run(target_floder, 40)
"""
def extract_python_code(prompt_str):
    import re
    python_code = re.findall(r"```python(.*?)```", prompt_str, re.DOTALL)[0]
    return python_code
# step1 将生成部分包装成函数
def gen_seed_generator_code(feature_result):
    # 生成python代码

    code_gen_prompt_str = f"""根据下面的jxl格式要求，写一个Python脚本，用来生成满足下面要求的jxl文件,要求要使用pillow_jxl库。
    要求：1.输出的文件夹要用户指定。
    2. 输出的函数接口名为run,里面要有输出文件夹、输出文件数量参数两个参数。
    3. 整个程序只有一个run函数，所有生成逻辑都要放到run函数中。
    4. 除了pillow_jxl库外，只能使用struct，numpy,os,sys库，严禁使用其他库。
    5. 不需要加入日志或者其他与生成jxl文件无关的功能。
    6. 不要加入任何调用run函数的示例代码。

    {feature_result}
    生成的代码是python代码：
    ```python
    //代码放到这里
    ```
    \\no_think
    """
    result = run_llm(code_gen_prompt_str)
    return result
def generate_featrue():
    # 生成特性
    feature_prompt_str = """根据你对jxl图像知识，列出jxl图像在文件格式上的10种算法特性。
    \\no_think
    """
    feature_result = run_llm(feature_prompt_str)
    return feature_result
    
# 先生成特性
feature_result = generate_featrue()

# step 2 设立死循环，不断尝试生成代码，直到成功为止
target_floder = "./j40/corpus"
rm_cmd =f"rm {target_floder}/*"
os.system(rm_cmd)
while True:
    rough_code = gen_seed_generator_code(feature_result)
    python_code = extract_python_code(rough_code)
    # 如果无法解析到python代码，则重新生成
    if not python_code: continue
    run_code = python_code + wrapper_code
    print(run_code)
    # 解析目标语法是否有问题
    try:
        ast.parse(run_code)
    except:
        print("\n代码语法错误\n")
        continue
    # 尝试运行目标代码，如果成功，则检查是否生成了种子
    try:
        exec(run_code)
        file_list = os.listdir(target_floder)
        if len(file_list) > 0:
            print("\n种子生成成功！\n")
            break
        else:
            print("\n生成的代码未能产生种子\n")
    except:
        print("\n生成代码运行错误！\n")
        continue